import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from JSAnimation import IPython_display

T = 20.
w = 2 * np.pi / T
ai = 1.  # Incident.
ar = 0.5  # Reflected.
d = 0
k = w**2 / 9.81  # Deep water wave.

t = np.arange(0, 100.5, 0.5)

def basic_animation(frames=100, interval=30):
    fig = plt.figure()
    ax = plt.axes(xlim=(0, 1000), ylim=(-2, 2))
    
    # Animated.
    line, = ax.plot([], [], 'b', lw=2)
    text = ax.text(1, 2.05, '')

    # Non-animated.
    x = np.arange(0, 1001)
    A = (ai**2 + ar**2 + 2*ai * ar * np.cos(2 * k * x + d))**(0.5)
    A1 = ai * np.cos(k * x) + ar * np.cos(k * x + d)
    A2 = ai * np.sin(k * x) - ar * np.sin(k * x + d)
    gamma = np.arctan2(A2, A1)
    ax.set_xlabel('Distance [m]')
    ax.set_ylabel('Surface Elevation [m]')
    ax.set_title('a_i = %s, a_r = %s, T = %s\nDeep Water Wave' % (ai, ar, T))
    ax.plot(x,  A, 'k:')
    ax.plot(x, -A, 'k:')
    ax.plot(x, [0]*len(x), 'k')
    line.set_data([], [])
    
    def init():
        return line, text

    def animate(i):
        text.set_text('Time [s] = %s' % t[k])
        y = A * np.cos(w * t[i] - gamma)
        line.set_data(x, y)
        return line, text

    return animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=frames, interval=interval)

basic_animation(frames=len(t))

# Wave 1.
k1, w1 = 5., 5.
c1 = w1 / k1

# Wave 2.
k2, w2 = 3., 4.5
c2 = w2 / k2

cg = (w2 - w1) / (k2 - k1)
cgc = (w2 + w1) / (k2 + k1)

deltg = 10 if cg <= 0 else 10.

def basic_animation(frames=200, interval=30, dt=0.1):
    fig = plt.figure(figsize=(16, 4))
    ax = plt.axes(xlim=(-15, 15), ylim=(-2, 2))
    
    # Animated.
    kw0 = dict(color='k', linestyle=':', alpha=0.5)
    kw1 = dict(alpha=0.5, linestyle='none', marker='o')
    text = ax.text([], 0, 'Cg')
    line0, = ax.plot([], [], color='gray')
    line1, = ax.plot([], [], 'r')
    line2, = ax.plot([], [], 'g')
    line3, = ax.plot([], [], **kw0)
    line4, = ax.plot([], [], **kw0)
    point0, = ax.plot([], [], 'r', **kw1)
    point1, = ax.plot([], [], 'g', **kw1)
    point2, = ax.plot([], [], 'gray', **kw1)

    # Non-animated.
    x = np.arange(-15, 15, 0.05)
    ax.set_xlabel('Distance')
    ax.set_ylabel('Amplitude')
    
    def init():
        return line0, line1, line2, line3, line4, point0, point1, point2, text

    def animate(k):
        t = k * dt
        eta1 = np.cos(k1 * x - w1 * t)
        eta2 = np.cos(k2 * x - w2 * t)
        etag = 2 * np.cos((k2 - k1) * x / 2. - (w2 - w1) * t / 2.)
        
        line0.set_data(x, eta1 + eta2)
        line1.set_data(x, eta1)
        line2.set_data(x, eta2)
        line3.set_data(x, etag)
        line4.set_data(x, -etag)
        point0.set_data(t * c1 - 10, 0)
        point1.set_data(t * c2 - 10, 0)
        point2.set_data(t * cgc - 10, 0)
        text.set_x(t * cg - deltg)
        return line0, line1, line2, line3, line4, point0, point1, point2, text

    return animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=frames, interval=interval)

dt = 0.1  # Time step
nstep = 150  # Number of time steps to take.
basic_animation(frames=60, dt=dt)

g = 9.81
denw = 1025.0  # Seawater density [kg/m**3].
sig = 7.3e-2  # Surface tension [N/m].
a = 1.0  # Wave amplitude [m].

L, h = 100.0, 50.0  # Wave height and water column depth.
k = 2 * np.pi / L
omega = np.sqrt((g * k + (sig / denw) * (k**3)) * np.tanh(k * h))
T = 2 * np.pi / omega
c = np.sqrt((g / k + (sig / denw) * k) * np.tanh(k * h))

# We'll solve the wave velocities in the `x` and `z` directions.
x, z = np.meshgrid(np.arange(0, 160, 10), np.arange(0, -80, -10),)
u, w = np.zeros_like(x), np.zeros_like(z)

def compute_vel(phase):
    u = a * omega * (np.cosh(k * (z+h)) / np.sinh(k*h)) * np.cos(k * x - phase)
    w = a * omega * (np.sinh(k * (z+h)) / np.sinh(k*h)) * np.sin(k * x - phase)
    mask = -z > h
    u[mask] = 0.0
    w[mask] = 0.0
    return u, w

def basic_animation(frames=91, interval=30, dt=0.3):
    fig = plt.figure(figsize=(8, 6))
    ax = plt.axes(xlim=(0, 150), ylim=(-70, 10))

    # Animated.
    quiver = ax.quiver(x, z, u, w, units='inches', scale=2)
    ax.quiverkey(quiver, 120, -60, 1,
                 label=r'1 m s$^{-1}$',
                 coordinates='data')
    line, = ax.plot([], [], 'b')

    # Non-animated.
    ax.plot([0, 150], [0, 0], 'k:')
    ax.set_ylabel('Depth [m]')
    ax.set_xlabel('Distance [m]')
    text = (r'$\lambda$ = %s m;  h = %s m;  kh = %2.3f;  h/L = %s' %
            (L, h, k * h, h/L))
    ax.text(10, -65, text)
    time_step = ax.text(10, -58, '')
    line.set_data([], [])

    def init():
        return line, quiver, time_step

    def animate(i):
        time = i * dt
        phase = omega * time
        eta = a * np.cos(x[0] * k - phase)
        u, w = compute_vel(phase)
        quiver.set_UVC(u, w)
        line.set_data(x[0], 5 * eta)
        time_step.set_text('Time = %s s' % time)
        return line, quiver, time_step

    return animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=frames, interval=interval)

dt = 0.3
basic_animation(dt=dt)

HTML(html)